package org.hong.monkey.scheduler

import java.util.Properties

import org.hong.monkey.scheduler.cluster.ExecutorInfo
import org.hong.monkey.util.{Utils, Distribution}
import org.hong.monkey.{Logging, TaskEndReason}
import org.hong.monkey.annotation.DeveloperApi
import org.hong.monkey.executor.TaskMetrics
import org.hong.monkey.storage.{BlockUpdatedInfo, BlockManagerId}

import scala.collection.Map
import scala.collection.mutable

@DeveloperApi
sealed trait MonkeyListenerEvent


@DeveloperApi
case class MonkeyListenerStageSubmitted(stageInfo: StageInfo, properties: Properties = null)
  extends MonkeyListenerEvent

@DeveloperApi
case class MonkeyListenerStageCompleted(stageInfo: StageInfo) extends MonkeyListenerEvent

@DeveloperApi
case class MonkeyListenerTaskStart(stageId: Int, stageAttemptId: Int, taskInfo: TaskInfo)
  extends MonkeyListenerEvent

@DeveloperApi
case class MonkeyListenerTaskGettingResult(taskInfo: TaskInfo) extends MonkeyListenerEvent

@DeveloperApi
case class MonkeyListenerTaskEnd(
    stageId: Int,
    stageAttemptId: Int,
    taskType: String,
    reason: TaskEndReason,
    taskInfo: TaskInfo,
    taskMetrics: TaskMetrics)
  extends MonkeyListenerEvent

@DeveloperApi
case class MonkeyListenerJobStart(
    jobId: Int,
    time: Long,
    stageInfos: Seq[StageInfo],
    properties: Properties = null)
  extends MonkeyListenerEvent {
  // Note: this is here for backwards-compatibility with older versions of this event which
  // only stored stageIds and not StageInfos:
  val stageIds: Seq[Int] = stageInfos.map(_.stageId)
}

@DeveloperApi
case class MonkeyListenerJobEnd(
    jobId: Int,
    time: Long,
    jobResult: JobResult)
  extends MonkeyListenerEvent

@DeveloperApi
case class MonkeyListenerEnvironmentUpdate(environmentDetails: Map[String, Seq[(String, String)]])
  extends MonkeyListenerEvent

@DeveloperApi
case class MonkeyListenerBlockManagerAdded(time: Long, blockManagerId: BlockManagerId, maxMem: Long)
  extends MonkeyListenerEvent

@DeveloperApi
case class MonkeyListenerBlockManagerRemoved(time: Long, blockManagerId: BlockManagerId)
  extends MonkeyListenerEvent

@DeveloperApi
case class MonkeyListenerUnpersistRDD(rddId: Int) extends MonkeyListenerEvent

@DeveloperApi
case class MonkeyListenerExecutorAdded(time: Long, executorId: String, executorInfo: ExecutorInfo)
  extends MonkeyListenerEvent

@DeveloperApi
case class MonkeyListenerExecutorRemoved(time: Long, executorId: String, reason: String)
  extends MonkeyListenerEvent

@DeveloperApi
case class MonkeyListenerBlockUpdated(blockUpdatedInfo: BlockUpdatedInfo) extends MonkeyListenerEvent
@DeveloperApi
case class MonkeyListenerExecutorMetricsUpdate(
    execId: String,
    taskMetrics: Seq[(Long, Int, Int, TaskMetrics)])
  extends MonkeyListenerEvent

@DeveloperApi
case class MonkeyListenerApplicationStart(
    appName: String,
    appId: Option[String],
    time: Long,
    monkeyUser: String,
    appAttemptId: Option[String],
    driverLogs: Option[Map[String, String]] = None) extends MonkeyListenerEvent

@DeveloperApi
case class MonkeyListenerApplicationEnd(time: Long) extends MonkeyListenerEvent

private[monkey] case class MonkeyListenerLogStart(monkeyVersion: String) extends MonkeyListenerEvent

@DeveloperApi
trait MonkeyListener {
  /**
   * Called when a stage completes successfully or fails, with information on the completed stage.
   */
  def onStageCompleted(stageCompleted: MonkeyListenerStageCompleted) { }

  /**
   * Called when a stage is submitted
   */
  def onStageSubmitted(stageSubmitted: MonkeyListenerStageSubmitted) { }

  /**
   * Called when a task starts
   */
  def onTaskStart(taskStart: MonkeyListenerTaskStart) { }

  /**
   * Called when a task begins remotely fetching its result (will not be called for tasks that do
   * not need to fetch the result remotely).
   */
  def onTaskGettingResult(taskGettingResult: MonkeyListenerTaskGettingResult) { }

  /**
   * Called when a task ends
   */
  def onTaskEnd(taskEnd: MonkeyListenerTaskEnd) { }

  /**
   * Called when a job starts
   */
  def onJobStart(jobStart: MonkeyListenerJobStart) { }

  /**
   * Called when a job ends
   */
  def onJobEnd(jobEnd: MonkeyListenerJobEnd) { }

  /**
   * Called when environment properties have been updated
   */
  def onEnvironmentUpdate(environmentUpdate: MonkeyListenerEnvironmentUpdate) { }

  /**
   * Called when a new block manager has joined
   */
  def onBlockManagerAdded(blockManagerAdded: MonkeyListenerBlockManagerAdded) { }

  /**
   * Called when an existing block manager has been removed
   */
  def onBlockManagerRemoved(blockManagerRemoved: MonkeyListenerBlockManagerRemoved) { }

  /**
   * Called when an RDD is manually unpersisted by the application
   */
  def onUnpersistRDD(unpersistRDD: MonkeyListenerUnpersistRDD) { }

  /**
   * Called when the application starts
   */
  def onApplicationStart(applicationStart: MonkeyListenerApplicationStart) { }

  /**
   * Called when the application ends
   */
  def onApplicationEnd(applicationEnd: MonkeyListenerApplicationEnd) { }

  /**
   * Called when the driver receives task metrics from an executor in a heartbeat.
   */
  def onExecutorMetricsUpdate(executorMetricsUpdate: MonkeyListenerExecutorMetricsUpdate) { }

  /**
   * Called when the driver registers a new executor.
   */
  def onExecutorAdded(executorAdded: MonkeyListenerExecutorAdded) { }

  /**
   * Called when the driver removes an executor.
   */
  def onExecutorRemoved(executorRemoved: MonkeyListenerExecutorRemoved) { }

  /**
   * Called when the driver receives a block update info.
   */
  def onBlockUpdated(blockUpdated: MonkeyListenerBlockUpdated) { }
}

@DeveloperApi
class StatsReportListener extends MonkeyListener with Logging {

  import org.hong.monkey.scheduler.StatsReportListener._

  private val taskInfoMetrics = mutable.Buffer[(TaskInfo, TaskMetrics)]()

  override def onTaskEnd(taskEnd: MonkeyListenerTaskEnd) {
    val info = taskEnd.taskInfo
    val metrics = taskEnd.taskMetrics
    if (info != null && metrics != null) {
      taskInfoMetrics += ((info, metrics))
    }
  }

  override def onStageCompleted(stageCompleted: MonkeyListenerStageCompleted) {
    implicit val sc = stageCompleted
    this.logInfo("Finished stage: " + stageCompleted.stageInfo)
    showMillisDistribution("task runtime:", (info, _) => Some(info.duration), taskInfoMetrics)

    showBytesDistribution("task result size:",
      (_, metric) => Some(metric.resultSize), taskInfoMetrics)

    // Runtime breakdown
    val runtimePcts = taskInfoMetrics.map { case (info, metrics) =>
      RuntimePercentage(info.duration, metrics)
    }
    showDistribution("executor (non-fetch) time pct: ",
      Distribution(runtimePcts.map(_.executorPct * 100)), "%2.0f %%")
    showDistribution("other time pct: ", Distribution(runtimePcts.map(_.other * 100)), "%2.0f %%")
    taskInfoMetrics.clear()
  }

}

private[monkey] object StatsReportListener extends Logging {

  // For profiling, the extremes are more interesting
  val percentiles = Array[Int](0, 5, 10, 25, 50, 75, 90, 95, 100)
  val probabilities = percentiles.map(_ / 100.0)
  val percentilesHeader = "\t" + percentiles.mkString("%\t") + "%"

  def extractDoubleDistribution(
      taskInfoMetrics: Seq[(TaskInfo, TaskMetrics)],
      getMetric: (TaskInfo, TaskMetrics) => Option[Double]): Option[Distribution] = {
    Distribution(taskInfoMetrics.flatMap { case (info, metric) => getMetric(info, metric) })
  }

  // Is there some way to setup the types that I can get rid of this completely?
  def extractLongDistribution(
      taskInfoMetrics: Seq[(TaskInfo, TaskMetrics)],
      getMetric: (TaskInfo, TaskMetrics) => Option[Long]): Option[Distribution] = {
    extractDoubleDistribution(
      taskInfoMetrics,
      (info, metric) => { getMetric(info, metric).map(_.toDouble) })
  }

  def showDistribution(heading: String, d: Distribution, formatNumber: Double => String) {
    val stats = d.statCounter
    val quantiles = d.getQuantiles(probabilities).map(formatNumber)
    logInfo(heading + stats)
    logInfo(percentilesHeader)
    logInfo("\t" + quantiles.mkString("\t"))
  }

  def showDistribution(
      heading: String,
      dOpt: Option[Distribution],
      formatNumber: Double => String) {
    dOpt.foreach { d => showDistribution(heading, d, formatNumber)}
  }

  def showDistribution(heading: String, dOpt: Option[Distribution], format: String) {
    def f(d: Double): String = format.format(d)
    showDistribution(heading, dOpt, f _)
  }

  def showDistribution(
      heading: String,
      format: String,
      getMetric: (TaskInfo, TaskMetrics) => Option[Double],
      taskInfoMetrics: Seq[(TaskInfo, TaskMetrics)]) {
    showDistribution(heading, extractDoubleDistribution(taskInfoMetrics, getMetric), format)
  }

  def showBytesDistribution(
      heading: String,
      getMetric: (TaskInfo, TaskMetrics) => Option[Long],
      taskInfoMetrics: Seq[(TaskInfo, TaskMetrics)]) {
    showBytesDistribution(heading, extractLongDistribution(taskInfoMetrics, getMetric))
  }

  def showBytesDistribution(heading: String, dOpt: Option[Distribution]) {
    dOpt.foreach { dist => showBytesDistribution(heading, dist) }
  }

  def showBytesDistribution(heading: String, dist: Distribution) {
    showDistribution(heading, dist, (d => Utils.bytesToString(d.toLong)): Double => String)
  }

  def showMillisDistribution(heading: String, dOpt: Option[Distribution]) {
    showDistribution(heading, dOpt,
      (d => StatsReportListener.millisToString(d.toLong)): Double => String)
  }

  def showMillisDistribution(
      heading: String,
      getMetric: (TaskInfo, TaskMetrics) => Option[Long],
      taskInfoMetrics: Seq[(TaskInfo, TaskMetrics)]) {
    showMillisDistribution(heading, extractLongDistribution(taskInfoMetrics, getMetric))
  }

  val seconds = 1000L
  val minutes = seconds * 60
  val hours = minutes * 60

  /**
   * Reformat a time interval in milliseconds to a prettier format for output
   */
  def millisToString(ms: Long): String = {
    val (size, units) =
      if (ms > hours) {
        (ms.toDouble / hours, "hours")
      } else if (ms > minutes) {
        (ms.toDouble / minutes, "min")
      } else if (ms > seconds) {
        (ms.toDouble / seconds, "s")
      } else {
        (ms.toDouble, "ms")
      }
    "%.1f %s".format(size, units)
  }
}

private case class RuntimePercentage(executorPct: Double, other: Double)

private object RuntimePercentage {
  def apply(totalTime: Long, metrics: TaskMetrics): RuntimePercentage = {
    val denom = totalTime.toDouble
    val exec = (metrics.executorRunTime) / denom
    val other = 1.0 - (exec)
    RuntimePercentage(exec, other)
  }
}